'''
Created on 2020年1月30日
获取所有股票数据
@author: JM
'''
import pandas as pd
import tushare as ts
from sqlalchemy import create_engine
import time
from app_config import get_pro, get_engine_ts


engine_ts = get_engine_ts()

counter = {'count': 0, "all": 0}

ts.set_token("09640b7c249de22f51c81fa10ebadcc2c31924864746e333e5ffbe2c")


def read_data():
    sql = """SELECT * FROM daily_basic_hfd"""
    df = pd.read_sql_query(sql, engine_ts)
    return df


def write_data(df_):
    res = df_.to_sql('daily_basic_hfd', engine_ts, index=False, if_exists='append', chunksize=5000)
    print(res)


def get_data(ts_code, start_date_str, end_date_str):
    df = ts.pro_bar(ts_code=ts_code, adj='hfq', start_date=start_date_str, end_date=end_date_str)
    return df


def get_daily_hdf(ts_code, start_date_str, end_date_str):
    counter['count'] += 1
    print(counter['count'], ": ", counter['all'])
    for _ in range(10):
        try:
            df = get_data(ts_code, start_date_str, end_date_str)
            # # 创建新的列 'id'，将 'ts_code' 和 'trade_date' 两列拼接起来
            df['id'] = df['ts_code'].astype(str) + df['trade_date'].astype(str)
            write_data(df)
        except Exception as e:
            print(f"An unexpected error occurred: {e}")
            time.sleep(10)
        else:
            return
    print("error error error error error" + ts_code)


if __name__ == '__main__':

    from _stock_basic import get_data as get_stock_basic

    # 获取20230501～20240430之间所有有交易的日期
    start_date_str = '20180101'
    # start_date_str = '20240101'
    end_date_str = '20240706'

    stock_basic = get_stock_basic()

    counter['all'] = len(stock_basic)

    for index, row in stock_basic.iterrows():  #
        # 每天的交易日期
        ts_code = row["ts_code"]
        print("股票" + str(ts_code))
        get_daily_hdf(ts_code, start_date_str, end_date_str)
